{
 "cells": [
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "%env LLM_API_KEY=sk-替换为自己的Qwen API Key\n",
    "%env LLM_BASE_URL=https://dashscope.aliyuncs.com/compatible-mode/v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:29.375334Z",
     "iopub.status.busy": "2024-10-12T02:26:29.374867Z",
     "iopub.status.idle": "2024-10-12T02:26:32.344846Z",
     "shell.execute_reply": "2024-10-12T02:26:32.344404Z",
     "shell.execute_reply.started": "2024-10-12T02:26:29.375310Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.10/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmExperimentalWarning: Using `tqdm.autonotebook.tqdm` in notebook mode. Use `tqdm.tqdm` instead to force console mode (e.g. in jupyter console)\n",
      "  from tqdm.autonotebook import tqdm, trange\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "langchain                     0.2.10\n",
      "langchain_core                0.2.28\n",
      "langchain_community           0.2.9\n",
      "pypdf                         4.3.1\n",
      "sentence_transformers         3.0.1\n",
      "chromadb                      0.5.4\n"
     ]
    }
   ],
   "source": [
    "import langchain, langchain_community, pypdf, sentence_transformers, chromadb, langchain_core\n",
    "\n",
    "for module in (langchain, langchain_core, langchain_community, pypdf, sentence_transformers, chromadb):\n",
    "    print(f\"{module.__name__:<30}{module.__version__}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:32.345483Z",
     "iopub.status.busy": "2024-10-12T02:26:32.345263Z",
     "iopub.status.idle": "2024-10-12T02:26:32.347562Z",
     "shell.execute_reply": "2024-10-12T02:26:32.347147Z",
     "shell.execute_reply.started": "2024-10-12T02:26:32.345468Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:32.348080Z",
     "iopub.status.busy": "2024-10-12T02:26:32.347961Z",
     "iopub.status.idle": "2024-10-12T02:26:32.360520Z",
     "shell.execute_reply": "2024-10-12T02:26:32.360096Z",
     "shell.execute_reply.started": "2024-10-12T02:26:32.348068Z"
    }
   },
   "outputs": [],
   "source": [
    "expr_version = 'retrieval_v10_contextual_compression'\n",
    "\n",
    "preprocess_output_dir = os.path.join(os.path.pardir, 'outputs', 'v1_20240713')\n",
    "expr_dir = os.path.join(os.path.pardir, 'experiments', expr_version)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 读取文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:32.361612Z",
     "iopub.status.busy": "2024-10-12T02:26:32.361486Z",
     "iopub.status.idle": "2024-10-12T02:26:33.883070Z",
     "shell.execute_reply": "2024-10-12T02:26:33.882597Z",
     "shell.execute_reply.started": "2024-10-12T02:26:32.361599Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "\n",
    "loader = PyPDFLoader(os.path.join(os.path.pardir, 'data', '2024全球经济金融展望报告.pdf'))\n",
    "documents = loader.load()\n",
    "\n",
    "qa_df = pd.read_excel(os.path.join(preprocess_output_dir, 'question_answer.xlsx'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.884033Z",
     "iopub.status.busy": "2024-10-12T02:26:33.883630Z",
     "iopub.status.idle": "2024-10-12T02:26:33.888186Z",
     "shell.execute_reply": "2024-10-12T02:26:33.887880Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.884020Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "53"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 文档切分"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.888735Z",
     "iopub.status.busy": "2024-10-12T02:26:33.888615Z",
     "iopub.status.idle": "2024-10-12T02:26:33.896633Z",
     "shell.execute_reply": "2024-10-12T02:26:33.896160Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.888723Z"
    }
   },
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "def split_docs(documents, filepath, chunk_size=400, chunk_overlap=40, seperators=['\\n\\n\\n', '\\n\\n'], force_split=False):\n",
    "    if os.path.exists(filepath) and not force_split:\n",
    "        print('found cache, restoring...')\n",
    "        return pickle.load(open(filepath, 'rb'))\n",
    "\n",
    "    splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size=chunk_size,\n",
    "        chunk_overlap=chunk_overlap,\n",
    "        separators=seperators\n",
    "    )\n",
    "    split_docs = splitter.split_documents(documents)\n",
    "    for chunk in split_docs:\n",
    "        chunk.metadata['uuid'] = str(uuid4())\n",
    "\n",
    "    pickle.dump(split_docs, open(filepath, 'wb'))\n",
    "\n",
    "    return split_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.897197Z",
     "iopub.status.busy": "2024-10-12T02:26:33.897076Z",
     "iopub.status.idle": "2024-10-12T02:26:33.902782Z",
     "shell.execute_reply": "2024-10-12T02:26:33.902361Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.897185Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found cache, restoring...\n"
     ]
    }
   ],
   "source": [
    "splitted_docs = split_docs(documents, os.path.join(preprocess_output_dir, 'split_docs.pkl'), chunk_size=500, chunk_overlap=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 检索"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "添加一个函数便于打印检索信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.903487Z",
     "iopub.status.busy": "2024-10-12T02:26:33.903340Z",
     "iopub.status.idle": "2024-10-12T02:26:33.907689Z",
     "shell.execute_reply": "2024-10-12T02:26:33.907246Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.903475Z"
    }
   },
   "outputs": [],
   "source": [
    "def pretty_print_docs(docs):\n",
    "    print(\n",
    "        f\"\\n{'-' * 100}\\n\".join(\n",
    "            [f\"Document {i + 1}:\" + (f\"\\nsimilarity: {d.state['query_similarity_score']:.4f}\" if hasattr(d, 'state') else '') + \"\\n\\n\" + d.page_content for i, d in enumerate(docs)]\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备基础检索器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.908546Z",
     "iopub.status.busy": "2024-10-12T02:26:33.908280Z",
     "iopub.status.idle": "2024-10-12T02:26:33.924794Z",
     "shell.execute_reply": "2024-10-12T02:26:33.924362Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.908533Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: cuda\n"
     ]
    }
   ],
   "source": [
    "import shutil\n",
    "import torch\n",
    "\n",
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "from langchain_community.vectorstores import Chroma\n",
    "\n",
    "model_path = 'BAAI/bge-large-zh-v1.5'\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "print(f'device: {device}')\n",
    "\n",
    "def get_embeddings(model_path):\n",
    "    embeddings = HuggingFaceBgeEmbeddings(\n",
    "        model_name=model_path,\n",
    "        model_kwargs={'device': device},\n",
    "        encode_kwargs={'normalize_embeddings': True},\n",
    "        query_instruction='为这个句子生成表示以用于检索相关文章：'\n",
    "    )\n",
    "    return embeddings\n",
    "\n",
    "def get_vector_db(embeddings, docs, db_name):\n",
    "    persist_directory = os.path.join(expr_dir, 'chroma', db_name)\n",
    "    shutil.rmtree(persist_directory, ignore_errors=True)\n",
    "\n",
    "    vector_db = Chroma.from_documents(\n",
    "        splitted_docs,\n",
    "        embedding=embeddings,\n",
    "        persist_directory=persist_directory\n",
    "    )\n",
    "    return vector_db"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:33.925384Z",
     "iopub.status.busy": "2024-10-12T02:26:33.925229Z",
     "iopub.status.idle": "2024-10-12T02:26:48.095392Z",
     "shell.execute_reply": "2024-10-12T02:26:48.093185Z",
     "shell.execute_reply.started": "2024-10-12T02:26:33.925372Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "embeddings = get_embeddings(model_path)\n",
    "vector_db = get_vector_db(embeddings, splitted_docs, 'contextual_compression')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "search_kwargs可以不指定，但本系列文章统一使用上下文长度为3进行对比，因此此处也设置这个参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:48.098382Z",
     "iopub.status.busy": "2024-10-12T02:26:48.097663Z",
     "iopub.status.idle": "2024-10-12T02:26:48.108321Z",
     "shell.execute_reply": "2024-10-12T02:26:48.105924Z",
     "shell.execute_reply.started": "2024-10-12T02:26:48.098308Z"
    }
   },
   "outputs": [],
   "source": [
    "base_retriever = vector_db.as_retriever(search_kwargs={'k': 3})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:48.111473Z",
     "iopub.status.busy": "2024-10-12T02:26:48.110635Z",
     "iopub.status.idle": "2024-10-12T02:26:48.172710Z",
     "shell.execute_reply": "2024-10-12T02:26:48.172182Z",
     "shell.execute_reply.started": "2024-10-12T02:26:48.111398Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 2 2024年\n",
      "图1：全球GDP增速（%）\n",
      "资料来源：IMF，中国银行研究院\n",
      "分区域看，全球经济复苏不均衡，各国存在较大差异。发达经济体增速明\n",
      "显放缓，预计2023年增速较2022年下降1个百分点。其中，欧元区和英国经\n",
      "济增速大幅下降，美国表现好于其他发达经济体。2023年三季度，欧元区和英\n",
      "国GDP环比增速均由之前的正增长转为负增长，分别下降0.1%和0.03%；美\n",
      "国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。新兴经济体增\n",
      "速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。其中，\n",
      "东南亚等出口型经济体增长承压，拉美、非洲等大宗商品出口国增速放缓，中\n",
      "东欧国家经济增速加快（图2）。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 9 2024年\n",
      "综合上述因素，“复苏+分化”将成为2024年全球经济主线逻辑，预计全\n",
      "年经济增速在2.5%左右，增速比2023年低0.2个百分点。\n",
      "（二）重点国别/地区经济形势研判\n",
      "1.美国经济增速将回落，高利率对经济的滞后影响将逐步显现\n",
      "2023年前三季度美国经济未受加息明显冲击，出现超预期增长。一、二季\n",
      "度GDP环比增长折年率分别为2.2%和2.1%，处于潜在增速附近，三季度大幅\n",
      "提高至4.9%，为2022年以来最好季度表现（图8）。其中，私人消费和政府财\n",
      "政支出是支撑美国经济增长的重要驱动力，前三季度上述两分项对美国经济增\n",
      "长的贡献率分别为64.4%和23.5%。后疫情时代消费者在服务领域的压抑需求\n",
      "得到释放。服务支出成为消费者支出领域中最稳定和增幅最大的板块，前三季\n",
      "度服务消费贡献了私人消费增长的58.7%。这主要得益于美国就业市场的稳健，\n",
      "1-10月非农就业人数均值远高于历史水平，失业率平均值为3.6%，低于2022\n",
      "年同期和国会预算办公室估算的非周期性失业率水平。美国劳动者实际收入保\n",
      "持增长，个人实际可支配收入月同比增速均值达4.2%，家庭财务状况良好，家\n",
      "庭财富对个人支出的正面效应高于历史水平。财政支出对经济增长的带动作用\n",
      "上升。2023年美国政府推进实施《通胀削减法案》《芯片与科学法案》，加大\n",
      "对清洁能源、半导体等产业的财政支持力度，前三季度政府支出平均每季度拉\n",
      "动GDP环比增长折年率0.7个百分点，而2022年该数值为-0.1。私人投资触底\n",
      "反弹，对经济的提振作用逐渐增强。受美国积极的产业政策和产业链回流等因\n",
      "素影响，制造业投资在二季度之前保持稳定增长，很大程度上对冲了利率抬升\n",
      "对美国房地产投资的负面影响，弱化了加息对私人投资的冲击。三季度虽然制\n",
      "造业投资开始转弱，但房地产投资的快速恢复又再度稳定了私人投资形势。前\n",
      "三季度私人投资对美国GDP环比增长折年率的季度拉动率均值为0.2%，高于\n",
      "上年同期的-0.8%。净出口对美国经济的拖累效应减弱。2023年美元维持强势\n",
      "表现，前三季度货物进口支出下降6.0%，而全球需求疲软导致美国货物出口规\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 7 2024年\n",
      "（图7）。10月，世贸组织将2023年全球货物贸易增速预测值下调0.9个百分\n",
      "点至0.8%，2023年国际贸易增长或为近几年最低水平。但近期东亚、东南亚等\n",
      "主要经济体出口下行趋势收窄，贸易呈现企稳迹象。9-10月，越南出口结束连\n",
      "续10个月的负增长态势，同比分别增长5.0%和6.7%。10月，韩国出口同比增\n",
      "长5.1%，是自2022年10月以来首次正增长。\n",
      "图7：全球货物贸易量指数和货物贸易价格指数（2010年=100）\n",
      "资料来源：荷兰经济分析局，中国银行研究院\n",
      "展望2024年，预计全球经济复苏将依旧疲软，且分化可能加剧。\n",
      "一是发达经济体面临的“遏通胀和稳增长”两难问题更加突出，货币政策\n",
      "走势将分化。疫后服务业复苏加快、劳动力市场紧平衡、地缘冲突加剧等因素\n",
      "导致发达经济体核心通胀黏性较强，通胀回落速度较慢。2023年10月，IMF\n",
      "将2024年全球通胀预测值由7月的5.2%上调至5.8%，并预测到2025年前大\n",
      "多数国家的通胀率都将高于央行目标。通胀反复风险使当前各主要央行对通胀\n",
      "走势保持警惕，预计在一段时间内政策利率将继续维持在高位。但随着经济下\n",
      "行压力加大，主要经济体央行货币政策走势将逐渐分化。截至11月中旬，在\n",
      "Macrobond跟踪的79家央行中，近1/5正处于降息周期。预计美联储大概率在\n",
      "2024年上半年停止加息，日本央行在通胀压力下可能于2024年年初退出负利\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(base_retriever.invoke('2023年全球经济增长的特点是什么？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用`LLMChainExtractor`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:48.174742Z",
     "iopub.status.busy": "2024-10-12T02:26:48.174598Z",
     "iopub.status.idle": "2024-10-12T02:26:48.363885Z",
     "shell.execute_reply": "2024-10-12T02:26:48.363410Z",
     "shell.execute_reply.started": "2024-10-12T02:26:48.174728Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain_community.chat_models import ChatOllama\n",
    "from langchain.retrievers import ContextualCompressionRetriever\n",
    "from langchain.retrievers.document_compressors import LLMChainExtractor\n",
    "\n",
    "llm = ChatOllama(base_url='http://localhost:11434', model='qwen2:7b-instruct')\n",
    "extractor_compressor = LLMChainExtractor.from_llm(llm)\n",
    "extractor_compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=extractor_compressor, base_retriever=base_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:26:48.364529Z",
     "iopub.status.busy": "2024-10-12T02:26:48.364362Z",
     "iopub.status.idle": "2024-10-12T02:27:04.077462Z",
     "shell.execute_reply": "2024-10-12T02:27:04.076939Z",
     "shell.execute_reply.started": "2024-10-12T02:26:48.364516Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "全球经济增长的特点如下：\n",
      "\n",
      "- 全球经济复苏不均衡，各国存在较大差异。\n",
      "- 发达经济体增速明显放缓，预计2023年增速较2022年下降1个百分点。其中，欧元区和英国经济增速大幅下降；美国表现好于其他发达经济体。\n",
      "- 欧元区和英国在2023年三季度的GDP环比增速均由之前的正增长转为负增长。\n",
      "- 美国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。\n",
      "- 新兴经济体增速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。其中，东南亚等出口型经济体增长承压；拉美、非洲等大宗商品出口国增速放缓；中东欧国家经济增速加快。\n",
      "\n",
      "以上是根据给定的上下文提取出的相关信息，用于回答关于2023年全球经济增长特点的问题。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "复苏+分化将成为2024年全球经济主线逻辑，预计全年经济增速在2.5%左右。美国经济增速将回落，高利率对经济的滞后影响将逐步显现。后疫情时代消费者在服务领域的压抑需求得到释放，服务支出成为消费者支出领域中最稳定和增幅最大的板块。美国就业市场稳健，非农就业人数均值远高于历史水平，失业率平均值为3.6%，个人实际可支配收入月同比增速均值达4.2%。财政支出对经济增长的带动作用上升。制造业投资在二季度之前保持稳定增长，很大程度上对冲了利率抬升对美国房地产投资的负面影响。净出口对美国经济的拖累效应减弱，2023年美元维持强势表现，前三季度货物进口支出下降6.0%，而全球需求疲软导致美国货物出口规模减少。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "\n",
      "近期东亚、东南亚等主要经济体出口下行趋势收窄，贸易呈现企稳迹象。9-10月，越南出口结束连续10个月的负增长态势，同比分别增长5.0%和6.7%。10月，韩国出口同比增长5.1%，是自2022年10月以来首次正增长。\n",
      "\n",
      "2024年预计全球经济复苏将依旧疲软，且分化可能加剧。发达经济体面临的“遏通胀和稳增长”两难问题更加突出，货币政策走势将分化。主要经济体央行在经济下行压力下可能会逐渐调整政策利率，其中美联储在2024年上半年停止加息的可能性较大。\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(extractor_compression_retriever.invoke('2023年全球经济增长的特点是什么？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用`LLMChainFilter`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`LLMChainFilter`只是要保留哪个文档，过滤掉哪个文档，不会更改文档本身的内容"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:04.078155Z",
     "iopub.status.busy": "2024-10-12T02:27:04.077945Z",
     "iopub.status.idle": "2024-10-12T02:27:04.081002Z",
     "shell.execute_reply": "2024-10-12T02:27:04.080622Z",
     "shell.execute_reply.started": "2024-10-12T02:27:04.078141Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import LLMChainFilter\n",
    "\n",
    "_filter = LLMChainFilter.from_llm(llm)\n",
    "filter_compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=_filter, base_retriever=base_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:04.081617Z",
     "iopub.status.busy": "2024-10-12T02:27:04.081483Z",
     "iopub.status.idle": "2024-10-12T02:27:05.913469Z",
     "shell.execute_reply": "2024-10-12T02:27:05.913017Z",
     "shell.execute_reply.started": "2024-10-12T02:27:04.081604Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 2 2024年\n",
      "图1：全球GDP增速（%）\n",
      "资料来源：IMF，中国银行研究院\n",
      "分区域看，全球经济复苏不均衡，各国存在较大差异。发达经济体增速明\n",
      "显放缓，预计2023年增速较2022年下降1个百分点。其中，欧元区和英国经\n",
      "济增速大幅下降，美国表现好于其他发达经济体。2023年三季度，欧元区和英\n",
      "国GDP环比增速均由之前的正增长转为负增长，分别下降0.1%和0.03%；美\n",
      "国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。新兴经济体增\n",
      "速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。其中，\n",
      "东南亚等出口型经济体增长承压，拉美、非洲等大宗商品出口国增速放缓，中\n",
      "东欧国家经济增速加快（图2）。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 7 2024年\n",
      "（图7）。10月，世贸组织将2023年全球货物贸易增速预测值下调0.9个百分\n",
      "点至0.8%，2023年国际贸易增长或为近几年最低水平。但近期东亚、东南亚等\n",
      "主要经济体出口下行趋势收窄，贸易呈现企稳迹象。9-10月，越南出口结束连\n",
      "续10个月的负增长态势，同比分别增长5.0%和6.7%。10月，韩国出口同比增\n",
      "长5.1%，是自2022年10月以来首次正增长。\n",
      "图7：全球货物贸易量指数和货物贸易价格指数（2010年=100）\n",
      "资料来源：荷兰经济分析局，中国银行研究院\n",
      "展望2024年，预计全球经济复苏将依旧疲软，且分化可能加剧。\n",
      "一是发达经济体面临的“遏通胀和稳增长”两难问题更加突出，货币政策\n",
      "走势将分化。疫后服务业复苏加快、劳动力市场紧平衡、地缘冲突加剧等因素\n",
      "导致发达经济体核心通胀黏性较强，通胀回落速度较慢。2023年10月，IMF\n",
      "将2024年全球通胀预测值由7月的5.2%上调至5.8%，并预测到2025年前大\n",
      "多数国家的通胀率都将高于央行目标。通胀反复风险使当前各主要央行对通胀\n",
      "走势保持警惕，预计在一段时间内政策利率将继续维持在高位。但随着经济下\n",
      "行压力加大，主要经济体央行货币政策走势将逐渐分化。截至11月中旬，在\n",
      "Macrobond跟踪的79家央行中，近1/5正处于降息周期。预计美联储大概率在\n",
      "2024年上半年停止加息，日本央行在通胀压力下可能于2024年年初退出负利\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(filter_compression_retriever.invoke('2023年全球经济增长的特点是什么？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用`LLMListwiseRerank`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`LLMListwiseRerank`使用zero-shot的方式对文档进行reranking，它的作用与`LLMChainFilter`类似，但是更加鲁邦也会耗费更多资源。建议使用强大一些的模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:05.914238Z",
     "iopub.status.busy": "2024-10-12T02:27:05.913990Z",
     "iopub.status.idle": "2024-10-12T02:27:05.938629Z",
     "shell.execute_reply": "2024-10-12T02:27:05.938174Z",
     "shell.execute_reply.started": "2024-10-12T02:27:05.914224Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import LLMListwiseRerank\n",
    "\n",
    "# 此处官方文档中介绍需要强大一些的LLM\n",
    "# https://python.langchain.com/docs/how_to/contextual_compression/\n",
    "\n",
    "# from langchain_openai import ChatOpenAI\n",
    "# gpt4o_llm = ChatOpenAI(model=\"gpt-4o-mini\", temperature=0)\n",
    "# _filter = LLMListwiseRerank.from_llm(gpt4o_llm, top_n=1)\n",
    "\n",
    "# 本文为了行文统一，依然还是使用Ollama中的Qwen2-7B\n",
    "# 此处需要用到函数调用功能，直接使用ChatOllama类的对象会报错，需要使用OllamaFunctions，但由于涉及到函数调用，实际使用时并不稳定\n",
    "\n",
    "from langchain_experimental.llms.ollama_functions import OllamaFunctions\n",
    "ollama_func_llm = OllamaFunctions(model='qwen2:7b-instruct', base_url='http://localhost:11434', format='json')\n",
    "\n",
    "_filter = LLMListwiseRerank.from_llm(ollama_func_llm, top_n=1)\n",
    "listwise_rerank_compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=_filter, base_retriever=base_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:05.939305Z",
     "iopub.status.busy": "2024-10-12T02:27:05.939182Z",
     "iopub.status.idle": "2024-10-12T02:27:08.973053Z",
     "shell.execute_reply": "2024-10-12T02:27:08.972589Z",
     "shell.execute_reply.started": "2024-10-12T02:27:05.939292Z"
    }
   },
   "outputs": [],
   "source": [
    "retrieved_docs = listwise_rerank_compression_retriever.invoke('2023年全球经济增长的特点是什么？')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:08.973769Z",
     "iopub.status.busy": "2024-10-12T02:27:08.973623Z",
     "iopub.status.idle": "2024-10-12T02:27:08.976477Z",
     "shell.execute_reply": "2024-10-12T02:27:08.975999Z",
     "shell.execute_reply.started": "2024-10-12T02:27:08.973755Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 2 2024年\n",
      "图1：全球GDP增速（%）\n",
      "资料来源：IMF，中国银行研究院\n",
      "分区域看，全球经济复苏不均衡，各国存在较大差异。发达经济体增速明\n",
      "显放缓，预计2023年增速较2022年下降1个百分点。其中，欧元区和英国经\n",
      "济增速大幅下降，美国表现好于其他发达经济体。2023年三季度，欧元区和英\n",
      "国GDP环比增速均由之前的正增长转为负增长，分别下降0.1%和0.03%；美\n",
      "国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。新兴经济体增\n",
      "速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。其中，\n",
      "东南亚等出口型经济体增长承压，拉美、非洲等大宗商品出口国增速放缓，中\n",
      "东欧国家经济增速加快（图2）。\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(retrieved_docs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用`EmbeddingFilter`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:08.977204Z",
     "iopub.status.busy": "2024-10-12T02:27:08.977016Z",
     "iopub.status.idle": "2024-10-12T02:27:08.987428Z",
     "shell.execute_reply": "2024-10-12T02:27:08.987000Z",
     "shell.execute_reply.started": "2024-10-12T02:27:08.977190Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain.retrievers.document_compressors import EmbeddingsFilter\n",
    "\n",
    "embeddings_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.6)\n",
    "embeddings_filter_compression_retriever = ContextualCompressionRetriever(\n",
    "    base_compressor=embeddings_filter, base_retriever=base_retriever\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:08.988071Z",
     "iopub.status.busy": "2024-10-12T02:27:08.987934Z",
     "iopub.status.idle": "2024-10-12T02:27:09.201889Z",
     "shell.execute_reply": "2024-10-12T02:27:09.201372Z",
     "shell.execute_reply.started": "2024-10-12T02:27:08.988059Z"
    }
   },
   "outputs": [],
   "source": [
    "retrieved_docs = embeddings_filter_compression_retriever.invoke('2023年全球经济增长的特点是什么？')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:09.202505Z",
     "iopub.status.busy": "2024-10-12T02:27:09.202375Z",
     "iopub.status.idle": "2024-10-12T02:27:09.205037Z",
     "shell.execute_reply": "2024-10-12T02:27:09.204652Z",
     "shell.execute_reply.started": "2024-10-12T02:27:09.202492Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document 1:\n",
      "similarity: 0.6317\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 2 2024年\n",
      "图1：全球GDP增速（%）\n",
      "资料来源：IMF，中国银行研究院\n",
      "分区域看，全球经济复苏不均衡，各国存在较大差异。发达经济体增速明\n",
      "显放缓，预计2023年增速较2022年下降1个百分点。其中，欧元区和英国经\n",
      "济增速大幅下降，美国表现好于其他发达经济体。2023年三季度，欧元区和英\n",
      "国GDP环比增速均由之前的正增长转为负增长，分别下降0.1%和0.03%；美\n",
      "国GDP环比增长折年率为4.9%，比二季度增速高2.8个百分点。新兴经济体增\n",
      "速与2022年大致持平，预计2023年增速比2022年下降0.1个百分点。其中，\n",
      "东南亚等出口型经济体增长承压，拉美、非洲等大宗商品出口国增速放缓，中\n",
      "东欧国家经济增速加快（图2）。\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 2:\n",
      "similarity: 0.6217\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 9 2024年\n",
      "综合上述因素，“复苏+分化”将成为2024年全球经济主线逻辑，预计全\n",
      "年经济增速在2.5%左右，增速比2023年低0.2个百分点。\n",
      "（二）重点国别/地区经济形势研判\n",
      "1.美国经济增速将回落，高利率对经济的滞后影响将逐步显现\n",
      "2023年前三季度美国经济未受加息明显冲击，出现超预期增长。一、二季\n",
      "度GDP环比增长折年率分别为2.2%和2.1%，处于潜在增速附近，三季度大幅\n",
      "提高至4.9%，为2022年以来最好季度表现（图8）。其中，私人消费和政府财\n",
      "政支出是支撑美国经济增长的重要驱动力，前三季度上述两分项对美国经济增\n",
      "长的贡献率分别为64.4%和23.5%。后疫情时代消费者在服务领域的压抑需求\n",
      "得到释放。服务支出成为消费者支出领域中最稳定和增幅最大的板块，前三季\n",
      "度服务消费贡献了私人消费增长的58.7%。这主要得益于美国就业市场的稳健，\n",
      "1-10月非农就业人数均值远高于历史水平，失业率平均值为3.6%，低于2022\n",
      "年同期和国会预算办公室估算的非周期性失业率水平。美国劳动者实际收入保\n",
      "持增长，个人实际可支配收入月同比增速均值达4.2%，家庭财务状况良好，家\n",
      "庭财富对个人支出的正面效应高于历史水平。财政支出对经济增长的带动作用\n",
      "上升。2023年美国政府推进实施《通胀削减法案》《芯片与科学法案》，加大\n",
      "对清洁能源、半导体等产业的财政支持力度，前三季度政府支出平均每季度拉\n",
      "动GDP环比增长折年率0.7个百分点，而2022年该数值为-0.1。私人投资触底\n",
      "反弹，对经济的提振作用逐渐增强。受美国积极的产业政策和产业链回流等因\n",
      "素影响，制造业投资在二季度之前保持稳定增长，很大程度上对冲了利率抬升\n",
      "对美国房地产投资的负面影响，弱化了加息对私人投资的冲击。三季度虽然制\n",
      "造业投资开始转弱，但房地产投资的快速恢复又再度稳定了私人投资形势。前\n",
      "三季度私人投资对美国GDP环比增长折年率的季度拉动率均值为0.2%，高于\n",
      "上年同期的-0.8%。净出口对美国经济的拖累效应减弱。2023年美元维持强势\n",
      "表现，前三季度货物进口支出下降6.0%，而全球需求疲软导致美国货物出口规\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Document 3:\n",
      "similarity: 0.6074\n",
      "\n",
      "全球经济金融展望报告\n",
      "中国银行研究院 7 2024年\n",
      "（图7）。10月，世贸组织将2023年全球货物贸易增速预测值下调0.9个百分\n",
      "点至0.8%，2023年国际贸易增长或为近几年最低水平。但近期东亚、东南亚等\n",
      "主要经济体出口下行趋势收窄，贸易呈现企稳迹象。9-10月，越南出口结束连\n",
      "续10个月的负增长态势，同比分别增长5.0%和6.7%。10月，韩国出口同比增\n",
      "长5.1%，是自2022年10月以来首次正增长。\n",
      "图7：全球货物贸易量指数和货物贸易价格指数（2010年=100）\n",
      "资料来源：荷兰经济分析局，中国银行研究院\n",
      "展望2024年，预计全球经济复苏将依旧疲软，且分化可能加剧。\n",
      "一是发达经济体面临的“遏通胀和稳增长”两难问题更加突出，货币政策\n",
      "走势将分化。疫后服务业复苏加快、劳动力市场紧平衡、地缘冲突加剧等因素\n",
      "导致发达经济体核心通胀黏性较强，通胀回落速度较慢。2023年10月，IMF\n",
      "将2024年全球通胀预测值由7月的5.2%上调至5.8%，并预测到2025年前大\n",
      "多数国家的通胀率都将高于央行目标。通胀反复风险使当前各主要央行对通胀\n",
      "走势保持警惕，预计在一段时间内政策利率将继续维持在高位。但随着经济下\n",
      "行压力加大，主要经济体央行货币政策走势将逐渐分化。截至11月中旬，在\n",
      "Macrobond跟踪的79家央行中，近1/5正处于降息周期。预计美联储大概率在\n",
      "2024年上半年停止加息，日本央行在通胀压力下可能于2024年年初退出负利\n"
     ]
    }
   ],
   "source": [
    "pretty_print_docs(retrieved_docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:09.205726Z",
     "iopub.status.busy": "2024-10-12T02:27:09.205570Z",
     "iopub.status.idle": "2024-10-12T02:27:09.218251Z",
     "shell.execute_reply": "2024-10-12T02:27:09.217791Z",
     "shell.execute_reply.started": "2024-10-12T02:27:09.205713Z"
    }
   },
   "outputs": [],
   "source": [
    "# pretty_print_docs(vector_db.as_retriever(search_type=\"similarity_score_threshold\", search_kwargs={'k': 2, 'score_threshold': 0.2}).invoke('2023年全球经济增长的特点是什么？'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 预测"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:09.218819Z",
     "iopub.status.busy": "2024-10-12T02:27:09.218669Z",
     "iopub.status.idle": "2024-10-12T02:27:09.224637Z",
     "shell.execute_reply": "2024-10-12T02:27:09.224186Z",
     "shell.execute_reply.started": "2024-10-12T02:27:09.218806Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain.llms import Ollama\n",
    "\n",
    "ollama_llm = Ollama(\n",
    "    model='qwen2:7b-instruct',\n",
    "    base_url='http://localhost:11434'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:09.225177Z",
     "iopub.status.busy": "2024-10-12T02:27:09.225055Z",
     "iopub.status.idle": "2024-10-12T02:27:09.232358Z",
     "shell.execute_reply": "2024-10-12T02:27:09.230066Z",
     "shell.execute_reply.started": "2024-10-12T02:27:09.225164Z"
    }
   },
   "outputs": [],
   "source": [
    "def rag(llm, retriever, question):\n",
    "    prompt_tmpl = \"\"\"\n",
    "你是一个金融分析师，擅长根据所获取的信息片段，对问题进行分析和推理。\n",
    "你的任务是根据所获取的信息片段（<<<<context>>><<<</context>>>之间的内容）回答问题。\n",
    "回答保持简洁，不必重复问题，不要添加描述性解释和与答案无关的任何内容。\n",
    "已知信息：\n",
    "<<<<context>>>\n",
    "{{knowledge}}\n",
    "<<<</context>>>\n",
    "\n",
    "问题：{{question}}\n",
    "请回答：\n",
    "\"\"\".strip()\n",
    "\n",
    "    chunks = retriever.invoke(question)\n",
    "    prompt = prompt_tmpl.replace('{{knowledge}}', '\\n\\n'.join([doc.page_content for doc in chunks])).replace('{{question}}', question)\n",
    "\n",
    "    return llm.invoke(prompt), chunks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:27:09.235622Z",
     "iopub.status.busy": "2024-10-12T02:27:09.234867Z",
     "iopub.status.idle": "2024-10-12T02:27:24.069501Z",
     "shell.execute_reply": "2024-10-12T02:27:24.069040Z",
     "shell.execute_reply.started": "2024-10-12T02:27:09.235553Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2023年全球经济增长的特点表现为\"复苏+分化\"的主线逻辑，预计全年经济增速为2.5%，比2022年降低0.2个百分点。发达经济体如欧元区和英国增长速度显著放缓，而美国表现相对较好。新兴经济体增长速度与前一年大致持平，但整体上仍面临一定压力。\n"
     ]
    }
   ],
   "source": [
    "print(rag(ollama_llm, extractor_compression_retriever, '2023年全球经济增长的特点是什么？')[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:30:04.747121Z",
     "iopub.status.busy": "2024-10-12T02:30:04.746300Z",
     "iopub.status.idle": "2024-10-12T02:30:04.770156Z",
     "shell.execute_reply": "2024-10-12T02:30:04.767966Z",
     "shell.execute_reply.started": "2024-10-12T02:30:04.747043Z"
    }
   },
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "\n",
    "prediction_df = qa_df[qa_df['dataset'] == 'test'][['uuid', 'question', 'qa_type', 'answer']].rename(columns={'answer': 'ref_answer'})\n",
    "\n",
    "def predict(llm, retriever, prediction_df):\n",
    "    prediction_df = prediction_df.copy()\n",
    "    answer_dict = {}\n",
    "\n",
    "    for idx, row in tqdm(prediction_df.iterrows(), total=len(prediction_df)):\n",
    "        uuid = row['uuid']\n",
    "        question = row['question']\n",
    "        answer, chunks = rag(ollama_llm, retriever, question)\n",
    "\n",
    "        answer_dict[question] = {\n",
    "            'uuid': uuid,\n",
    "            'ref_answer': row['ref_answer'],\n",
    "            'gen_answer': answer,\n",
    "            'chunks': chunks\n",
    "        }\n",
    "    prediction_df.loc[:, 'gen_answer'] = prediction_df['question'].apply(lambda q: answer_dict[q]['gen_answer'])\n",
    "    prediction_df.loc[:, 'chunks'] = prediction_df['question'].apply(lambda q: answer_dict[q]['chunks'])\n",
    "\n",
    "    return prediction_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:30:08.769047Z",
     "iopub.status.busy": "2024-10-12T02:30:08.768263Z",
     "iopub.status.idle": "2024-10-12T02:30:08.777093Z",
     "shell.execute_reply": "2024-10-12T02:30:08.775134Z",
     "shell.execute_reply.started": "2024-10-12T02:30:08.768972Z"
    }
   },
   "outputs": [],
   "source": [
    "pred_dfs = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T02:30:08.781892Z",
     "iopub.status.busy": "2024-10-12T02:30:08.781158Z",
     "iopub.status.idle": "2024-10-12T02:55:58.114871Z",
     "shell.execute_reply": "2024-10-12T02:55:58.114453Z",
     "shell.execute_reply.started": "2024-10-12T02:30:08.781818Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicting with Vanilla Retriever:\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b988a5d9c09c45dd8af6f7e63517f4ca",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicting with LLMChainExtractor:\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39b3f0fb760044948837879b5c78643a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicting with LLMChainFilter:\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8ec4490ab3a74db4ab1e120c5888b993",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "predicting with EmbeddingsFilter:\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5b4f35053094d0c9dff21f3b18a1e2b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for key, retriever in zip(\n",
    "    ['Vanilla Retriever', 'LLMChainExtractor', 'LLMChainFilter', 'EmbeddingsFilter'],\n",
    "    [base_retriever, extractor_compression_retriever, filter_compression_retriever, embeddings_filter_compression_retriever]\n",
    "):\n",
    "    if key in pred_dfs:\n",
    "        continue\n",
    "    print(f'predicting with {key}:')\n",
    "    pred_dfs[key] = predict(ollama_func_llm, retriever, prediction_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 评估"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T03:39:15.011221Z",
     "iopub.status.busy": "2024-10-12T03:39:15.010363Z",
     "iopub.status.idle": "2024-10-12T03:39:15.065567Z",
     "shell.execute_reply": "2024-10-12T03:39:15.064615Z",
     "shell.execute_reply.started": "2024-10-12T03:39:15.011145Z"
    }
   },
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "import time\n",
    "\n",
    "judge_llm = ChatOpenAI(\n",
    "    api_key=os.environ['LLM_API_KEY'],\n",
    "    base_url=os.environ['LLM_BASE_URL'],\n",
    "    model_name='qwen2-72b-instruct',\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "def evaluate(judge_llm, prediction_df):\n",
    "    \"\"\"\n",
    "    对预测结果进行打分\n",
    "    :param prediction_df: 预测结果，需要包含问题，参考答案，生成的答案，列名分别为question, ref_answer, gen_answer\n",
    "    :return 打分模型原始返回结果\n",
    "    \"\"\"\n",
    "    prompt_tmpl = \"\"\"\n",
    "你是一个经济学博士，现在我有一系列问题，有一个助手已经对这些问题进行了回答，你需要参照参考答案，评价这个助手的回答是否正确，仅回复“是”或“否”即可，不要带其他描述性内容或无关信息。\n",
    "问题：\n",
    "<question>\n",
    "{{question}}\n",
    "</question>\n",
    "\n",
    "参考答案：\n",
    "<ref_answer>\n",
    "{{ref_answer}}\n",
    "</ref_answer>\n",
    "\n",
    "助手回答：\n",
    "<gen_answer>\n",
    "{{gen_answer}}\n",
    "</gen_answer>\n",
    "请评价：\n",
    "    \"\"\"\n",
    "    results = []\n",
    "\n",
    "    for _, row in tqdm(prediction_df.iterrows(), total=len(prediction_df)):\n",
    "        question = row['question']\n",
    "        ref_answer = row['ref_answer']\n",
    "        gen_answer = row['gen_answer']\n",
    "\n",
    "        prompt = prompt_tmpl.replace('{{question}}', question).replace('{{ref_answer}}', str(ref_answer)).replace('{{gen_answer}}', gen_answer).strip()\n",
    "        result = judge_llm.invoke(prompt).content\n",
    "        results.append(result)\n",
    "\n",
    "        time.sleep(1)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T03:39:15.311067Z",
     "iopub.status.busy": "2024-10-12T03:39:15.310274Z",
     "iopub.status.idle": "2024-10-12T03:49:34.798032Z",
     "shell.execute_reply": "2024-10-12T03:49:34.797639Z",
     "shell.execute_reply.started": "2024-10-12T03:39:15.310992Z"
    }
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "08b11032efb645bf88f8cf5ed10c32ae",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla Retriever unique raw score: ['是' '否']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7a046f3bcd754602a75dc0c907fe57b2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLMChainExtractor unique raw score: ['是' '否']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b269d6ed561d493f958b68b67ec36a23",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLMChainFilter unique raw score: ['是' '否']\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b742776679a641eab2a89d03a264b558",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/100 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EmbeddingsFilter unique raw score: ['否' '是']\n"
     ]
    }
   ],
   "source": [
    "for key, pred_df in pred_dfs.items():\n",
    "    pred_df['raw_score'] = evaluate(judge_llm, pred_df)\n",
    "    pred_df['score'] = (pred_df['raw_score'] == '是').astype(int)\n",
    "    print(f\"{key} unique raw score: {pred_df['raw_score'].unique()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T03:49:34.798863Z",
     "iopub.status.busy": "2024-10-12T03:49:34.798724Z",
     "iopub.status.idle": "2024-10-12T03:49:34.801388Z",
     "shell.execute_reply": "2024-10-12T03:49:34.801065Z",
     "shell.execute_reply.started": "2024-10-12T03:49:34.798849Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython import display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T03:49:34.803440Z",
     "iopub.status.busy": "2024-10-12T03:49:34.803266Z",
     "iopub.status.idle": "2024-10-12T03:49:34.818780Z",
     "shell.execute_reply": "2024-10-12T03:49:34.818421Z",
     "shell.execute_reply.started": "2024-10-12T03:49:34.803427Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Vanilla Retriever accuracy: 0.7\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qa_type</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>detailed</td>\n",
       "      <td>0.688172</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>large_context</td>\n",
       "      <td>0.857143</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLMChainExtractor accuracy: 0.6\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qa_type</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>detailed</td>\n",
       "      <td>0.602151</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>large_context</td>\n",
       "      <td>0.571429</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLMChainFilter accuracy: 0.49\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qa_type</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>detailed</td>\n",
       "      <td>0.483871</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>large_context</td>\n",
       "      <td>0.571429</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EmbeddingsFilter accuracy: 0.14\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>qa_type</th>\n",
       "      <th>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>detailed</td>\n",
       "      <td>0.118280</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>large_context</td>\n",
       "      <td>0.428571</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for key, pred_df in pred_dfs.items():\n",
    "    pred_df['score'] = (pred_df['raw_score'] == '是').astype(int)\n",
    "    print(f\"{key} accuracy: {pred_df['score'].mean()}\")\n",
    "\n",
    "    display.display_html(pred_df[['qa_type', 'score']].groupby('qa_type').mean().reset_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-10-12T03:49:34.819441Z",
     "iopub.status.busy": "2024-10-12T03:49:34.819311Z",
     "iopub.status.idle": "2024-10-12T03:49:35.284562Z",
     "shell.execute_reply": "2024-10-12T03:49:35.284096Z",
     "shell.execute_reply.started": "2024-10-12T03:49:34.819428Z"
    }
   },
   "outputs": [],
   "source": [
    "for key, pred_df in pred_dfs.items():\n",
    "    pred_df.to_excel(os.path.join(expr_dir, f'{expr_version}_{key}_prediction.xlsx'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
